package com.itaming.lycheeframework.processor;

import com.google.common.collect.TreeMultimap;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.Filer;
import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import javax.lang.model.type.DeclaredType;
import javax.lang.model.type.TypeMirror;
import javax.lang.model.util.ElementFilter;
import javax.lang.model.util.SimpleAnnotationValueVisitor8;
import javax.tools.Diagnostic;
import javax.tools.FileObject;
import javax.tools.StandardLocation;
import java.io.*;
import java.lang.annotation.Annotation;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.stream.Collectors;

/**
 * SPI注解处理器
 *
 * @author A.Ming
 */
public class SpiProcessor extends AbstractProcessor {

    /**
     * SPI资源文件路径
     */
    private static final String RESOURCE_FILE_PATH = "META-INF/services/";

    /**
     * 收集SPI实现的有序映射
     */
    private final TreeMultimap<String, String> implementerMap = TreeMultimap.create();

    /**
     * 支持的注解类型
     *
     * @return 注解类型集合
     */
    @Override
    public Set<String> getSupportedAnnotationTypes() {
        return Set.of(SpiImpl.class.getName());
    }

    /**
     * 支持的版本
     *
     * @return SourceVersion
     */
    @Override
    public SourceVersion getSupportedSourceVersion() {
        return SourceVersion.latestSupported();
    }

    /**
     * 处理注解
     *
     * @param annotations 注解类型集合
     * @param roundEnv    当前编译轮次的环境信息
     */
    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        return !roundEnv.processingOver() ? collectImplementers(annotations, roundEnv) : writeImplementers();
    }

    /**
     * 收集实现
     *
     * @param annotations 注解类型集合
     * @param roundEnv    当前编译轮次的环境信息
     * @return boolean
     */
    private boolean collectImplementers(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        // 获取所有标注了 @SpiImpl 的类
        Set<? extends Element> elements = roundEnv.getElementsAnnotatedWith(SpiImpl.class);

        // 收集实现
        for (Element element : elements) {
            if (element.getKind() != ElementKind.CLASS) {
                continue;
            }

            TypeElement implementer = (TypeElement) element;

            AnnotationMirror annotationMirror = getAnnotationMirror(element, SpiImpl.class);
            if (annotationMirror == null) {
                continue;
            }

            Set<TypeMirror> typeMirrors = getValueTypes(annotationMirror);
            if (typeMirrors.isEmpty()) {
                continue;
            }

            for (TypeMirror typeMirror : typeMirrors) {
                String interfaceName = getTypeQualifiedName(typeMirror);
                Name implementerName = implementer.getQualifiedName();

                if (checkImplementer(implementer.asType(), typeMirror)) {
                    implementerMap.put(interfaceName, getTypeQualifiedName(implementer));
                } else {
                    String message = implementerName + " does not implement " + interfaceName;
                    processingEnv.getMessager().printMessage(Diagnostic.Kind.ERROR, message, element, annotationMirror);
                    return false;
                }
            }

        }

        return true;
    }

    /**
     * 写入实现
     *
     * @return boolean
     */
    private boolean writeImplementers() {
        NavigableMap<String, Collection<String>> map = implementerMap.asMap();
        Filer filer = processingEnv.getFiler();

        for (Map.Entry<String, Collection<String>> entry : map.entrySet()) {
            String interfaceName = entry.getKey();
            Collection<String> implementerNames = entry.getValue();
            String fileName = RESOURCE_FILE_PATH + interfaceName;
            Set<String> toAdd = new TreeSet<>(String::compareTo);

            // 读取已存在的文件，否则手动编写的内容会丢失
            try {
                FileObject fileObject = filer.getResource(StandardLocation.CLASS_OUTPUT, "", fileName);
                toAdd.addAll(readServiceFile(fileObject));
            } catch (IOException ignored) {
                // 文件不存在，忽略
            }

            toAdd.addAll(implementerNames);
            try {
                FileObject fileObject = filer.createResource(StandardLocation.CLASS_OUTPUT, "", fileName);
                writeServiceFile(fileObject, toAdd);
            } catch (IOException e) {
                processingEnv.getMessager().printMessage(Diagnostic.Kind.ERROR, "Failed to write service file: " + fileName + ", error: " + e.getMessage());
                return false;
            }
        }

        return true;
    }

    /**
     * 读取服务文件
     *
     * @param fileObject 文件对象
     * @return 已存在的实现类名集合
     * @throws IOException
     */
    private List<String> readServiceFile(FileObject fileObject) throws IOException {
        List<String> serviceList = new ArrayList<>();
        try (
            InputStreamReader isr = new InputStreamReader(fileObject.openInputStream(), StandardCharsets.UTF_8);
            BufferedReader reader = new BufferedReader(isr)
        ) {
            String line;
            while ((line = reader.readLine()) != null) {
                int commentIndex = line.indexOf('#');
                if (commentIndex >= 0) {
                    line = line.substring(0, commentIndex);
                }
                line = line.strip();
                if (!line.isEmpty()) {
                    serviceList.add(line);
                }
            }
            return serviceList;
        }
    }

    /**
     * 写入服务文件
     *
     * @param fileObject       文件对象
     * @param implementerNames 实现类名集合
     * @throws IOException
     */
    private void writeServiceFile(FileObject fileObject, Collection<String> implementerNames) throws IOException {
        try (
            OutputStreamWriter osw = new OutputStreamWriter(fileObject.openOutputStream(), StandardCharsets.UTF_8);
            BufferedWriter writer = new BufferedWriter(osw)
        ) {
            for (String implementerName : implementerNames) {
                writer.write(implementerName);
                writer.newLine();
            }
        }
    }

    /**
     * 获取注解
     *
     * @param element         实现类
     * @param annotationClass 注解类
     * @return 注解
     */
    private AnnotationMirror getAnnotationMirror(Element element, Class<? extends Annotation> annotationClass) {
        String annotationClassName = annotationClass.getCanonicalName();
        for (AnnotationMirror annotationMirror : element.getAnnotationMirrors()) {
            String name = getTypeQualifiedName(annotationMirror);
            if (name.contentEquals(annotationClassName)) {
                return annotationMirror;
            }
        }
        return null;
    }

    /**
     * 读取注解的<code>value</code>元素的类型集合
     *
     * @param annotationMirror 注解
     * @return 类型集合
     */
    private Set<TypeMirror> getValueTypes(AnnotationMirror annotationMirror) {
        return getAnnotationValue(annotationMirror, "value")
            .accept(new SimpleAnnotationValueVisitor8<Set<TypeMirror>, Void>() {
                @Override
                public Set<TypeMirror> visitType(TypeMirror typeMirror, Void v) {
                    return Collections.singleton(typeMirror);
                }

                @Override
                public Set<TypeMirror> visitArray(List<? extends AnnotationValue> values, Void v) {
                    return values.stream()
                        .flatMap(value -> value.accept(this, null).stream())
                        .collect(Collectors.toSet());
                }
            }, null);
    }

    /**
     * 获取注解元素的值
     *
     * @param annotationMirror 注解
     * @param elementName      注解元素名
     * @return 注解元素的值
     */
    public AnnotationValue getAnnotationValue(AnnotationMirror annotationMirror, String elementName) {
        Objects.requireNonNull(annotationMirror);
        Objects.requireNonNull(elementName);
        for (Map.Entry<ExecutableElement, AnnotationValue> entry : getAnnotationValuesWithDefaults(annotationMirror).entrySet()) {
            if (entry.getKey().getSimpleName().contentEquals(elementName)) {
                return entry.getValue();
            }
        }
        String name = getTypeQualifiedName(annotationMirror);
        throw new IllegalStateException(String.format("@%s does not define an element %s()", name, elementName));
    }

    /**
     * 获取注解元素的值映射
     *
     * @param annotation 注解
     * @return 注解元素的值映射
     */
    public Map<ExecutableElement, AnnotationValue> getAnnotationValuesWithDefaults(AnnotationMirror annotation) {
        Map<ExecutableElement, AnnotationValue> values = new HashMap<>(32);
        Map<? extends ExecutableElement, ? extends AnnotationValue> declaredValues = annotation.getElementValues();

        for (ExecutableElement method : ElementFilter.methodsIn(annotation.getAnnotationType().asElement().getEnclosedElements())) {
            if (declaredValues.containsKey(method)) {
                values.put(method, declaredValues.get(method));
            } else if (method.getDefaultValue() != null) {
                values.put(method, method.getDefaultValue());
            } else {
                String name = getTypeQualifiedName(method);
                throw new IllegalStateException("Unset annotation value without default should never happen: " + name + '.' + method.getSimpleName() + "()");
            }
        }

        return Collections.unmodifiableMap(values);
    }

    /**
     * 检查实现类是否实现了接口
     *
     * @param implementerType 实现类类型
     * @param interfaceType   接口类型
     * @return boolean
     */
    private boolean checkImplementer(TypeMirror implementerType, TypeMirror interfaceType) {
        return processingEnv.getTypeUtils().isSubtype(implementerType, interfaceType);
    }

    /**
     * 获取类型完整限定名
     *
     * @param element 实现类
     * @return 类型完整限定名
     */
    private String getTypeQualifiedName(Element element) {
        return getTypeQualifiedName(element != null ? element.asType() : null);
    }

    /**
     * 获取类型完整限定名
     *
     * @param annotation 注解
     * @return 类型完整限定名
     */
    private String getTypeQualifiedName(AnnotationMirror annotation) {
        return getTypeQualifiedName(annotation != null ? annotation.getAnnotationType() : null);
    }

    /**
     * 获取类型完整限定名完整限定名
     *
     * @param type 类型
     * @return 类型完整限定名
     */
    private String getTypeQualifiedName(TypeMirror type) {
        if (type == null) {
            return null;
        }
        if (type instanceof DeclaredType declaredType) {
            Element enclosingElement = declaredType.asElement().getEnclosingElement();
            if (enclosingElement instanceof TypeElement) {
                return getQualifiedName(enclosingElement) + "$" + declaredType.asElement().getSimpleName().toString();
            } else {
                return getQualifiedName(declaredType.asElement());
            }
        }
        return type.toString();
    }

    /**
     * 获取完整限定名
     *
     * @param element 实现类
     * @return 完整限定名
     */
    private String getQualifiedName(Element element) {
        if (element instanceof QualifiedNameable) {
            return ((QualifiedNameable) element).getQualifiedName().toString();
        }
        return element.toString();
    }

}
